from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta
from typing import Literal
from urllib.parse import quote

import pymongo
import pytz
from loguru import logger
from sqlalchemy import Column, Integer, String, and_, create_engine, delete, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session
from tqdm import tqdm

mongo_url = "mongodb://rwuser:Gr#54WE24KL!dfjh5!!lkfg874@localhost:27017/bigdata?authSource=admin&directConnection=true"
mongo_client = pymongo.MongoClient(mongo_url, minPoolSize=5)
bigdata = mongo_client.bigdata
clip_annotated_record = bigdata.clip_annotated_record
clip_col = bigdata.clip
check_frame_label_col = bigdata.check_frame_label
gt_task_item_col = bigdata.annotated_ground_truth_task_item

mysql_pwd = quote("zdrive@Test2023")
mysql_stage_engine = create_engine(
    f"mysql+pymysql://root:{mysql_pwd}@172.25.117.3:3306/bdc_v3?charset=utf8mb4"
)
mysql_engine = create_engine(
    "mysql+pymysql://root:Hl76##fds87FS093498!aac667@localhost:3306/bdc?charset=utf8mb4"
)


Base = declarative_base()

CntByClientType = dict[Literal["sensetime", "imotion", "zdrive"], int]
PeriodType = Literal["day", "week", "month", "quarter", "year"]
ProjectType = Literal["driving", "parking"]


CLIENT_LIST = [
    dict(name="sensetime", zh_name="商汤"),
    dict(name="imotion", zh_name="知行"),
    dict(name="zdrive", zh_name="大卓"),
]
POOL_SIZE = 10


def has_parking_tag(tags: dict[str, list[str]]):
    route_tags = tags.get("route") or []
    return "indoor_parking" in route_tags or "open_parking" in route_tags


def has_driving_tag(tags: dict[str, list[str]]):
    # TODO 再确认一下这个逻辑
    route_tags = tags.get("route") or []
    return "highway" in route_tags or (
        "city_road" in route_tags
        and "open_parking" not in route_tags
        and "indoor_parking" not in route_tags
    )


# 行车和泊车对应的标注类型
# 23 的话需要 tag 满足 route-city_road
# TODO 需求量先写死
driving_annotate_types_dict = {
    2: {
        "name": "行车 车道线",
        "demand_clip_cnt": 75_0000,
        "valid_tag": lambda tags: True,
    },
    5: {
        "name": "行车 红绿灯",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    15: {
        "name": "行车 PVB-城区",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    16: {
        "name": "行车 PVB-高速",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    22: {
        "name": "行车 OCC",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    23: {"name": "行车 GOP", "demand_clip_cnt": 50_0000, "valid_tag": has_driving_tag},
    24: {
        "name": "行车 相机遮挡",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_driving_tag,
    },
    25: {
        "name": "行车 环境识别",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_driving_tag,
    },
    27: {
        "name": "行车 交通标志",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    28: {
        "name": "行车 光斑",
        "demand_clip_cnt": 30_0000,
        "valid_tag": lambda tags: True,
    },
    32: {
        "name": "行车 3D-OCC",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
}
driving_annotate_types = list(driving_annotate_types_dict.keys())


# 23 的话需要 tag 满足 route-indoor_parking route-open_parking
parking_annotate_types_dict = {
    17: {
        "name": "泊车 PVB",
        "demand_clip_cnt": 100_0000,
        "valid_tag": lambda tags: True,
    },
    19: {
        "name": "泊车 车位",
        "demand_clip_cnt": 100_0000,
        "valid_tag": lambda tags: True,
    },
    20: {
        "name": "泊车 路面标识",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    21: {
        "name": "泊车 OCC",
        "demand_clip_cnt": 100_0000,
        "valid_tag": lambda tags: True,
    },
    23: {"name": "泊车 GOP", "demand_clip_cnt": 50_0000, "valid_tag": has_parking_tag},
    24: {
        "name": "泊车 相机遮挡",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_parking_tag,
    },
    25: {
        "name": "泊车 环境识别",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_parking_tag,
    },
    31: {
        "name": "泊车 车道线",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
}

parking_annotate_types = list(parking_annotate_types_dict.keys())

annotate_types_by_project = {
    "parking": parking_annotate_types_dict,
    "driving": driving_annotate_types_dict,
}

SH_TZ = pytz.timezone("Asia/Shanghai")


class LabelDashboardORM(Base):
    __tablename__ = "label_dashboard"

    id = Column(Integer, primary_key=True, autoincrement=True)
    project = Column(String(255), nullable=False, comment="项目名称")
    period = Column(
        String(10),
        nullable=False,
        comment="统计周期",
    )
    annotate_type = Column(Integer, nullable=False, comment="标注类型")
    annotate_name = Column(String(255), nullable=False, comment="标注类型名称")
    start_date = Column(String(10), nullable=False, comment="开始日期")
    end_date = Column(String(10), nullable=False, comment="结束日期")

    demand_clip_cnt = Column(Integer, default=0, comment="需求 clip 数")
    # demand_frame_cnt = Column(Integer, default=0, comment="需求 frame 数")
    curated_clip_cnt = Column(Integer, default=0, comment="已送标 clip 数")
    curated_frame_cnt = Column(Integer, default=0, comment="已送标 frame 数")
    labeled_clip_cnt = Column(Integer, default=0, comment="已标注 clip 数")
    labeled_frame_cnt = Column(Integer, default=0, comment="已标注 frame 数")
    post_processed_clip_cnt = Column(Integer, default=0, comment="后处理 clip 数")
    post_processed_frame_cnt = Column(Integer, default=0, comment="后处理 frame 数")
    accept_finished_clip_cnt = Column(Integer, default=0, comment="已验收 clip 数")
    accept_finished_frame_cnt = Column(Integer, default=0, comment="已验收 frame 数")
    accept_passed_frame_cnt = Column(Integer, default=0, comment="验收通过帧数")
    sensetime_delivered_clip_cnt = Column(
        Integer, default=0, comment="商汤交付 clip 数"
    )
    sensetime_delivered_frame_cnt = Column(
        Integer, default=0, comment="商汤交付 frame 数"
    )
    imotion_delivered_clip_cnt = Column(Integer, default=0, comment="知行交付 clip 数")
    imotion_delivered_frame_cnt = Column(
        Integer, default=0, comment="知行交付 frame 数"
    )
    zdrive_delivered_clip_cnt = Column(Integer, default=0, comment="大卓交付 clip 数")
    zdrive_delivered_frame_cnt = Column(Integer, default=0, comment="大卓交付 frame 数")


def handle_curated_fields(
    dashboard: LabelDashboardORM,
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    start_ts = int(datetime.strptime(start_date, "%Y-%m-%d").timestamp() * 1000)
    end_ts = int(datetime.strptime(end_date, "%Y-%m-%d").timestamp() * 1000)
    clip_info_map = {}
    clip_obj_map = {}
    with mysql_engine.connect() as conn:
        sql = text(
            """
            select id from send_label_task_tab 
            where ctime >= :start_ts and ctime < :end_ts 
            and annotate_type = :annotate_type
            """
        ).bindparams(start_ts=start_ts, end_ts=end_ts, annotate_type=annotate_type)

        all_task = conn.execute(sql).fetchall()
        if len(all_task) == 0:
            return

        for (task_id,) in all_task:
            sql = text(
                """
                select clip_id, count(distinct frame_id)
                from send_label_record_tab 
                where send_label_task_id = :task_id
                group by clip_id
                """
            ).bindparams(task_id=task_id)

            clip_infos = conn.execute(sql).fetchall()
            clip_ids = tuple([clip_id for clip_id, _ in clip_infos])

            clip_info_map.update(
                {
                    clip_id: dict(clip_id=clip_id, frame_cnt=frame_cnt)
                    for clip_id, frame_cnt in clip_infos
                }
            )
            clip_objs = clip_col.find(
                {"_id": {"$in": clip_ids}},
                {"tags": 1},
            )
            clip_obj_map.update({clip_obj["_id"]: clip_obj for clip_obj in clip_objs})

    def _accumulate_curate_info(clip_obj):
        dashboard.curated_frame_cnt += clip_obj["frame_cnt"]
        dashboard.curated_clip_cnt += 1

    for clip_id in clip_info_map:
        if annotate_types_by_project[project][annotate_type]["valid_tag"](
            clip_obj_map[clip_id]["tags"]
        ):
            _accumulate_curate_info(clip_info_map[clip_id])


def handle_labeled_fields(
    dashboard: LabelDashboardORM,
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    start_ts = int(datetime.strptime(start_date, "%Y-%m-%d").timestamp() * 1000)
    end_ts = int(datetime.strptime(end_date, "%Y-%m-%d").timestamp() * 1000)
    filter = dict(
        confirmed_time={"$gte": start_ts, "$lt": end_ts},
        # state={"$in": [2, 4]},
        state=2,  # 确认
        annotate_type=annotate_type,
    )
    # TODO 这里的 tag 和 clip 的 tags 是同步的吗？
    projection = dict(
        clip_id=1,
        state=1,
        annotate_type=1,
        confirmed_time=1,
        frame_count=1,
        tags=1,
        is_vis_processed=1,
    )
    clip_objs = clip_annotated_record.find(filter, projection).to_list()

    def _accumulate_label_info(clip_obj):
        dashboard.labeled_frame_cnt += clip_obj["frame_count"]
        dashboard.labeled_clip_cnt += 1
        dashboard.post_processed_clip_cnt += (
            1 if clip_obj.get("is_vis_processed", False) else 0
        )
        dashboard.post_processed_frame_cnt += (
            clip_obj["frame_count"] if clip_obj.get("is_vis_processed", False) else 0
        )

    for clip_obj in clip_objs:
        if annotate_types_by_project[project][annotate_type]["valid_tag"](
            clip_obj["tags"]
        ):
            _accumulate_label_info(clip_obj)


def handle_accept_fields(
    dashboard: LabelDashboardORM,
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    # mongodb 存储的是 utc 时间，start_date 和 end_date 是 shanghai 时区，
    # 所以这里要转换成带 shanghai 时区的时间，才能正确查询
    start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").astimezone(SH_TZ)
    end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").astimezone(SH_TZ)
    filter = dict(
        accept_finished_at={"$gte": start_date_obj, "$lt": end_date_obj},
        accept_status=20,
        annotate_type=annotate_type,
    )
    # TODO 这里的 tag 和 clip 的 tags 是同步的吗？
    projection = dict(
        clip_id=1,
        annotate_type=1,
        tags=1,
        accept_finished_at=1,
    )
    clip_records = clip_annotated_record.find(filter, projection).to_list()
    clip_ids = [clip_obj["clip_id"] for clip_obj in clip_records]

    # 从 check_frame_col 表中查询用户提交过的帧数，并按照 judging_standard 进行分组
    frame_standard_groups = check_frame_label_col.aggregate(
        [
            {"$match": {"clip_id": {"$in": clip_ids}, "annotate_type": annotate_type}},
            {
                "$group": {
                    "_id": {
                        "judging_standard": "$judging_standard",
                        "clip_id": "$clip_id",
                    },
                    "frame_count": {"$sum": 1},
                }
            },
        ]
    ).to_list()
    clip_standard_frame_cnt_map = defaultdict(dict)
    for group in frame_standard_groups:
        clip_id = group["_id"]["clip_id"]
        standard = group["_id"]["judging_standard"]
        clip_standard_frame_cnt_map[clip_id][standard] = group["frame_count"]

    def _accumulate_accept_info(clip_obj):
        clip_id = clip_obj["clip_id"]
        dashboard.accept_finished_frame_cnt += sum(
            list(clip_standard_frame_cnt_map[clip_id].values())
        )
        dashboard.accept_passed_frame_cnt += sum(
            [
                clip_standard_frame_cnt_map[clip_id].get("A", 0),
                clip_standard_frame_cnt_map[clip_id].get("B", 0),
            ]
        )
        dashboard.accept_finished_clip_cnt += 1

    for clip_obj in clip_records:
        if annotate_types_by_project[project][annotate_type]["valid_tag"](
            clip_obj["tags"]
        ):
            _accumulate_accept_info(clip_obj)


def handle_delivery_fields(
    dashboard: LabelDashboardORM,
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    filter = dict(
        annotate_type=annotate_type,
        status="success",
        # updated_at={"$gte": start_date, "$lt": end_date},
    )
    # NOTE 历史原因，updated_at 同时存在 string 和 datetime 类型，所以这里要用 $or
    filter.update(
        {
            "$or": [
                {"updated_at": {"$gte": start_date, "$lt": end_date}},
                {
                    "updated_at": {
                        "$gte": datetime.strptime(start_date, "%Y-%m-%d"),
                        "$lt": datetime.strptime(end_date, "%Y-%m-%d"),
                    }
                },
            ]
        }
    )

    projection = dict(
        clip_id=1,
        annotate_type=1,
        status=1,
        updated_at=1,
        metadata=1,
        supplier=1,
        num_frames=1,
    )

    # clip id 可能重复，要进行去重
    gt_task_item_objs = (
        gt_task_item_col.find(filter, projection).sort({"updated_at": -1}).to_list()
    )
    clip_ids = [obj["clip_id"] for obj in gt_task_item_objs]

    clip_info_map = {
        clip["_id"]: clip
        for clip in clip_col.find(
            {"_id": {"$in": clip_ids}}, {"frame_valid_count": 1, "tags": 1}
        )
    }

    def _accumulate_delivery_info(task_item):
        client = task_item.get("supplier")
        frame_cnt = task_item.get("num_frames")
        if frame_cnt is None:
            if task_item.get("metadata") is None:
                frame_cnt = 0
            else:
                frame_cnt = task_item.get("metadata", {}).get("NUM_GT", 0)

        if client == "sensetime":
            dashboard.sensetime_delivered_clip_cnt += 1
            dashboard.sensetime_delivered_frame_cnt += int(frame_cnt)
        elif client == "imotion":
            dashboard.imotion_delivered_clip_cnt += 1
            dashboard.imotion_delivered_frame_cnt += int(frame_cnt)
        elif client == "zdrive":
            dashboard.zdrive_delivered_clip_cnt += 1
            dashboard.zdrive_delivered_frame_cnt += int(frame_cnt)

    clip_id_set = set()
    for task_item in gt_task_item_objs:
        clip_id = task_item["clip_id"]
        if clip_id in clip_id_set:
            continue
        clip_id_set.add(clip_id)
        tags = clip_info_map[clip_id]["tags"]
        # 如果是 23，则判断 tag 是否满足条件，其他情况直接累加
        if annotate_types_by_project[project][annotate_type]["valid_tag"](tags):
            _accumulate_delivery_info(task_item)


def handle_one_annotate_type(
    project: ProjectType,
    annotate_type: int,
    period: PeriodType,
    start_date_str: str,
    end_date_str: str,
):
    annotate_name = None
    demand_clip_cnt = 0
    if project == "driving":
        annotate_name = driving_annotate_types_dict[annotate_type]["name"]
        demand_clip_cnt = (
            driving_annotate_types_dict[annotate_type]["demand_clip_cnt"]
            if period == "total"
            else 0
        )
    if project == "parking":
        annotate_name = parking_annotate_types_dict[annotate_type]["name"]
        demand_clip_cnt = (
            parking_annotate_types_dict[annotate_type]["demand_clip_cnt"]
            if period == "total"
            else 0
        )
    dashboard = LabelDashboardORM(
        project=project,
        annotate_type=annotate_type,
        annotate_name=annotate_name,
        start_date=start_date_str,
        end_date=end_date_str,
        period=period,
        demand_clip_cnt=demand_clip_cnt,
        curated_clip_cnt=0,
        curated_frame_cnt=0,
        labeled_clip_cnt=0,
        labeled_frame_cnt=0,
        post_processed_clip_cnt=0,
        post_processed_frame_cnt=0,
        accept_finished_clip_cnt=0,
        accept_finished_frame_cnt=0,
        accept_passed_frame_cnt=0,
        sensetime_delivered_clip_cnt=0,
        sensetime_delivered_frame_cnt=0,
        imotion_delivered_clip_cnt=0,
        imotion_delivered_frame_cnt=0,
        zdrive_delivered_clip_cnt=0,
        zdrive_delivered_frame_cnt=0,
    )
    handle_curated_fields(
        dashboard, project, annotate_type, start_date_str, end_date_str
    )
    handle_labeled_fields(
        dashboard, project, annotate_type, start_date_str, end_date_str
    )
    handle_accept_fields(
        dashboard, project, annotate_type, start_date_str, end_date_str
    )
    handle_delivery_fields(
        dashboard, project, annotate_type, start_date_str, end_date_str
    )

    return dashboard


def build_dashboard_by_project(
    project: str, period: str, start_date: str, end_date: str
):
    logger.info(
        f"build {project} {period} dashboard for date range [{start_date}, {end_date})"
    )
    data_list = []
    thread_tasks = []
    if project == "parking":
        annotate_types = parking_annotate_types
    if project == "driving":
        annotate_types = driving_annotate_types
    with ThreadPoolExecutor(max_workers=POOL_SIZE) as executor:
        for annotate_type in annotate_types:
            thread_tasks.append(
                executor.submit(
                    handle_one_annotate_type,
                    project,
                    annotate_type,
                    period,
                    start_date,
                    end_date,
                )
            )

        for future in tqdm(as_completed(thread_tasks), total=len(thread_tasks)):
            dashboard = future.result()
            data_list.append(dashboard)
    return data_list


def build_dashboard(period: str, start_date: str, end_date: str):
    driving_data_list = build_dashboard_by_project(
        "driving", period, start_date, end_date
    )
    parking_data_list = build_dashboard_by_project(
        "parking", period, start_date, end_date
    )
    return [*driving_data_list, *parking_data_list]

    # with pd.ExcelWriter("label_dashboard.xlsx") as writer:
    #     # 将 DataFrame 写入不同的 sheet
    #     driving_df.to_excel(writer, sheet_name="行车", index=False)
    #     parking_df.to_excel(writer, sheet_name="泊车", index=False)


def insert_dashboard(dashboard_list: list[LabelDashboardORM]):
    with Session(mysql_stage_engine) as session:
        session.add_all(dashboard_list)
        session.commit()


def delete_dashboard(period: str, start_date: str):
    with Session(mysql_stage_engine) as session:
        del_stmt = delete(LabelDashboardORM).where(
            and_(
                LabelDashboardORM.period == period,
                LabelDashboardORM.start_date == start_date,
            )
        )
        ret = session.execute(del_stmt)
        session.commit()
    return ret.rowcount


def build_dashboard_by_date(date: str | None = None):
    if date is None:
        # 自动运行，计算的是前一天的数据
        end_date = datetime.now().astimezone(SH_TZ).date()
        start_date = end_date - timedelta(days=1)
    else:
        # 用户指定日期，则计算当天数据
        start_date = datetime.strptime(date, "%Y-%m-%d").astimezone(SH_TZ).date()
        end_date = start_date + timedelta(days=1)

    start_date_str, end_date_str = (
        start_date.strftime("%Y-%m-%d"),
        end_date.strftime("%Y-%m-%d"),
    )

    # 每天执行一次，计算前一天的数据
    data_list = build_dashboard("day", start_date_str, end_date_str)
    delete_row_cnt = delete_dashboard("day", start_date_str)
    logger.info(
        f"delete day dashboard for date range [{start_date_str}, {end_date_str}): {delete_row_cnt}"
    )
    insert_dashboard(data_list)
    logger.info(
        f"insert day dashboard for date range [{start_date_str}, {end_date_str}): {len(data_list)}"
    )

    # 每天执行一次，计算累计数据
    if end_date == datetime.now().astimezone(SH_TZ).date():
        start_date = (
            datetime.strptime("2023-01-01", "%Y-%m-%d").astimezone(SH_TZ).date()
        )
        start_date_str = start_date.strftime("%Y-%m-%d")

        data_list = build_dashboard("total", start_date_str, end_date_str)
        delete_row_cnt = delete_dashboard("total", start_date_str)
        logger.info(
            f"delete total dashboard for date range [{start_date_str}, {end_date_str}): {delete_row_cnt}"
        )
        insert_dashboard(data_list)
        logger.info(
            f"insert total dashboard for date range [{start_date_str}, {end_date_str}): {len(data_list)}"
        )

    # 每周一执行一次计算上周数据，或计算当前周数据
    if end_date == datetime.now().astimezone(SH_TZ).date() or end_date.weekday() == 0:
        start_date = end_date - timedelta(
            days=7 if end_date.weekday() == 0 else end_date.weekday()
        )
        start_date_str = start_date.strftime("%Y-%m-%d")
        data_list = build_dashboard("week", start_date_str, end_date_str)
        logger.info(
            f"delete week dashboard for date range[{start_date_str}, {end_date_str}): {delete_row_cnt}"
        )
        delete_row_cnt = delete_dashboard("week", start_date_str)
        insert_dashboard(data_list)
        logger.info(
            f"insert week dashboard for date range[{start_date_str}, {end_date_str}): {len(data_list)}"
        )


if __name__ == "__main__":
    build_dashboard_by_date()
    # start_date = "2025-01-01"
    # cur_date = datetime.now().strftime("%Y-%m-%d")
    # while cur_date >= start_date:
    #     build_dashboard_by_date(cur_date)
    #     cur_date = (
    #         datetime.strptime(cur_date, "%Y-%m-%d") - timedelta(days=1)
    #     ).strftime("%Y-%m-%d")
